# -*- coding: utf-8 -*-

# 导入pyspark
from pyspark import SparkContext
from pyspark.mllib.clustering import KMeans, KMeansModel
from pyspark.mllib.recommendation import ALS, Rating
from pyspark.mllib.linalg import Vectors

# 导入可视化
import matplotlib
matplotlib.use('Agg') # 不回显
import matplotlib.pyplot as plt

# 导入数据库操作
import MySQLdb

# 函数-操作数据库
def insertNewRecordForDecisionTreeRegression(sql):
	# 连接数据库
	try:
		conn = MySQLdb.connect(host='localhost',user='root',passwd='root',db='sparkDemo')
	except Exception, e:
		print e
	# cursor对象操作数据库
	cursor = conn.cursor()
	# 执行SQL
	try:
		cursor.execute(sql)
	except Exception, e:
		print e
	# 提交操作
	conn.commit()
	# 关闭指针
	cursor.close()
	# 关闭数据库
	conn.close()
	return 1

# 函数-类别
def genrePre(x):
	if len(x) == 2:
		return (int(x[1]), x[0])
	pass

def isGenre(arr):
	print("==============================================\n")
	retArr = [arr[1], arr[2]]
	i = 0
	for field in arr[5:24]:
		if(field == "1"):
			retArr.append(genre_records_filter[i][1])
		i = i + 1
	print str(retArr)
	return (int(arr[0]), retArr)


# main函数部分
sc = SparkContext("yarn-client", "K-Means Spark App")

# 加载数据集 - genre & item
genrePath = "hdfs://192.168.119.141:9100/data/movie/u.genre"
itemPath = "hdfs://192.168.119.141:9100/data/movie/u.item"
ratingPath = "hdfs://192.168.119.141:9100/data/movie/u.data"
# 数据集
genre_data_raw = sc.textFile(genrePath)
item_data_raw = sc.textFile(itemPath)
rating_data_raw = sc.textFile(ratingPath)
# 分隔
genre_records = genre_data_raw.map(lambda x: x.split("|"))
item_records = item_data_raw.map(lambda x: x.split("|"))
ratings_records = rating_data_raw.map(lambda x: x.split('\t'))

# 数据 - ALS rating
ratings_data = ratings_records.map(lambda x: Rating(int(x[0]), int(x[1]), float(x[2])))
# 数据 - movie + genre
genre_records_filter = genre_records.map(genrePre).collect()
# 缓存rdd
ratings_data.cache()

# 写入文件
f = open('getKMeansInfo.txt','w')

# train(ratings, rank, iterations=5, lambda_=0.01, blocks=-1, nonnegative=False, seed=None)[source]
# 返回两个RDD - userFeatures / productFeatures
alsModel = ALS.train(ratings_data, 50, 10, 0.1)
movieAls = alsModel.productFeatures()
# train/test 数据准备 
trainMovies, testMovies = movieAls.randomSplit([0.8, 0.2], 17)
# train(trainData, numClusters, 最大迭代次数, 最大轮回次数,)
movieClusterModel = KMeans.train(trainMovies.map(lambda x: Vectors.dense(x[1])), 6, maxIterations=10, 
									runs=3, initializationMode="random")
# 信息保存到数据库（wcss）
wcss = movieClusterModel.computeCost(testMovies.map(lambda x: Vectors.dense(x[1])))
sql = "insert into queue(name, status, mse) values('%s', %d, %f)"%("KMeansModel", 1, wcss)
insertNewRecordForDecisionTreeRegression(sql)

# 验证
movieCluster = testMovies.map(lambda x: (x[0], movieClusterModel.predict(Vectors.dense(x[1]))))
# 数据 - 用户电影(movieId, (userId, rating))
userMovie = ratings_records.map(lambda x: ( int(x[1]), (int(x[0]), float(x[2])) ))
# 数据 - 评分(movieId, avg(rating))
rating = movieAls.map(lambda x: (x[0], sum(x[1]) / len(x[1])))
# 数据 - movie (movieId, (catogory, avg(rating)))
item_records_filter = item_records.map(isGenre).join(rating)
item_records_filter.cache()

# 链接电影信息(movieId, (predictCluster, (catogory, avg(rating)))
movieLast = movieCluster.join(item_records_filter)

# f.write("\n<br>===TestMovies Predict Sample===<br>\n")
# f.write(str(movieLast.take(5)))

# 保存电影信息(predictCluster, (movieId, (catogory, avg(rating))))
movieSave = movieLast.map(lambda (x, y): (y[0], (x, y[1])))
# 根据clusterId分组map(predictCluster, [(movieId, (catogory, avg(rating))), (), ...])
moviefile = movieSave.groupByKey().mapValues(list)

# f.write("\n<br>===cluster Group Info===<br>\n")
# f.write(str(moviefile.take(5)))
# f.write("\n<br>===cluster end===<br>\n")

# 取6个cluster可视化[(movieId, (catogory, avg(rating))), (), ...]
clusterOne = moviefile.filter(lambda (x, y): x==0).map(lambda (x, y): y).collect()
clusterTwo = moviefile.filter(lambda (x, y): x==1).map(lambda (x, y): y).collect()
clusterThree = moviefile.filter(lambda (x, y): x==2).map(lambda (x, y): y).collect()
clusterFour = moviefile.filter(lambda (x, y): x==3).map(lambda (x, y): y).collect()
clusterFive = moviefile.filter(lambda (x, y): x==4).map(lambda (x, y): y).collect()
clusterSix = moviefile.filter(lambda (x, y): x==5).map(lambda (x, y): y).collect()

f.write("\n<br>===cluster01===<br>\n")
f.write(str(clusterOne))
f.write("\n<br>===cluster02===<br>\n")
f.write(str(clusterTwo))
f.write("\n<br>===cluster03===<br>\n")
f.write(str(clusterThree))
f.write("\n<br>===cluster04===<br>\n")
f.write(str(clusterFour))
f.write("\n<br>===cluster05===<br>\n")
f.write(str(clusterFive))
f.write("\n<br>===cluster06===<br>\n")
f.write(str(clusterSix))
f.write("\n<br>===cluster end===<br>\n")

#　准备 ｘ / ｙ 画图
arrayOneX = []
arrayOneY = []
arrayTwoX = []
arrayTwoY = []
arrayThreeX = []
arrayThreeY = []
arrayFourX = []
arrayFourY = []
arrayFiveX = []
arrayFiveY = []
arraySixX = []
arraySixY = []
arrayOne = clusterOne[0]
arrayTwo = clusterTwo[0]
arrayThree = clusterThree[0]
arrayFour = clusterFour[0]
arrayFive = clusterFive[0]
arraySix = clusterSix[0]


for arrEle in arrayOne:
	arrayOneX.append(arrEle[0])
	arrayOneY.append(arrEle[1][1])
for arrEle in arrayTwo:
	arrayTwoX.append(arrEle[0]+2000)
	arrayTwoY.append(arrEle[1][1])
for arrEle in arrayThree:
	arrayThreeX.append(arrEle[0]+4000)
	arrayThreeY.append(arrEle[1][1])
for arrEle in arrayFour:
	arrayFourX.append(arrEle[0]+6000)
	arrayFourY.append(arrEle[1][1])
for arrEle in arrayFive:
	arrayFiveX.append(arrEle[0]+8000)
	arrayFiveY.append(arrEle[1][1])
for arrEle in arraySix:
	arraySixX.append(arrEle[0]+10000)
	arraySixY.append(arrEle[1][1])

# 画图
plt.figure(num=1, figsize=(8,6))
plt.title('KMeans Cluster Model', size=14)
plt.xlabel('Movie', size=14)
plt.ylabel('Rating', size=14)

plt.plot(arrayOneX, arrayOneY, 'ro', label="cluster01")
plt.plot(arrayTwoX, arrayTwoY, 'bo', label="cluster02")
plt.plot(arrayThreeX, arrayThreeY, 'go', label="cluster03")
plt.plot(arrayFourX, arrayFourY, 'ko', label="cluster04")
plt.plot(arrayFiveX, arrayFiveY, 'co', label="cluster05")
plt.plot(arraySixX, arraySixY, 'mo', label="cluster06")


# K 点

clusterCen = movieClusterModel.clusterCenters
f.write(str(clusterCen))
f.write("====")

# cenX = []
# cenY = []
# for i in range(0, len(clusterCen)):
# 	cenX.append(clusterCen[i][0])
# 	cenY.append(clusterCen[i][1])

# for i in range(0, len(clusterCen)):
# 	arrShow = (cenX[i], cenY[i])
# 	textShow = str(arrShow)
# 	plt.text(cenX[i], cenY[i], textShow, color="red", fontsize=6)

# plt.plot(cenX, cenY, 'y+', label="ClusterCenter")


# 保存图片文件
plt.savefig('KMeans.png', format='png')

# 关闭文件指针
f.close()

# sc关闭
sc.stop()
